export class FlappyBirdEnv {
    constructor() {
        this.width = 800;
        this.height = 600;
        
        this.GRAVITY = 0.3;
        this.JUMP_FORCE = -3;
        this.PIPE_SPEED = 2;
        this.PIPE_SPAWN_INTERVAL = 180;

        // 确保立即生成第一个管道
        this.frameCount = this.PIPE_SPAWN_INTERVAL;
        this.reset();
    }

    reset() {
        this.bird = {
            x: 100,
            y: this.height / 2,
            velocity: 0,
            radius: 15
        };
        
        // 重置时生成一个初始管道
        this.pipes = [];
        this.spawnPipe();
        
        this.score = 0;
        this.gameOver = false;
        this.frameCount = this.PIPE_SPAWN_INTERVAL; // 确保马上生成下一个管道
        return this.getState();
    }

    step(action) {
        this.frameCount++;
        
        // Bird movement
        if (action === 1) { // Jump
            this.bird.velocity = this.JUMP_FORCE;
        }
        
        this.bird.velocity += this.GRAVITY;
        this.bird.y += this.bird.velocity;

        // 管道生成逻辑
        if (this.frameCount >= this.PIPE_SPAWN_INTERVAL) {
            this.spawnPipe();
            this.frameCount = 0;
        }

        // 更新管道位置
        for (let pipe of this.pipes) {
            pipe.x -= this.PIPE_SPEED;
        }
        
        this.pipes = this.pipes.filter(pipe => pipe.x > -pipe.width);

        // 检查碰撞
        if (this.checkCollision()) {
            this.gameOver = true;
            return [this.getState(), -20, true];  // 增加碰撞惩罚
        }

        // 计算奖励
        let reward = 0;
        
        // 1. 基于位置的奖励：鼓励小鸟保持在屏幕中间区域
        const centerY = this.height / 2;
        const distanceToCenter = Math.abs(this.bird.y - centerY);
        const positionReward = distanceToCenter < 100 ? 0.2 : -distanceToCenter / this.height;
        
        // 2. 基于速度的奖励：严厉惩罚过大的速度
        const velocityPenalty = Math.abs(this.bird.velocity) > 8 ? -1 : -Math.abs(this.bird.velocity) / 10;
        
        // 3. 通过管道的奖励
        let pipeReward = 0;
        const closestPipe = this.pipes.find(pipe => pipe.x + pipe.width >= this.bird.x);
        if (closestPipe) {
            // 当接近管道时给予额外奖励
            const horizontalDist = closestPipe.x - this.bird.x;
            if (horizontalDist < 100 && horizontalDist > -closestPipe.width) {
                if (this.bird.y > closestPipe.y && this.bird.y < closestPipe.y + closestPipe.gap) {
                    pipeReward = 0.5;  // 在管道间隙中时给予奖励
                }
            }
        }
        
        for (let pipe of this.pipes) {
            if (!pipe.passed && pipe.x + pipe.width < this.bird.x) {
                pipe.passed = true;
                this.score++;
                pipeReward = 10.0;  // 通过管道的奖励
            }
        }

        // 4. 接近边缘的惩罚
        const edgePenalty = (this.bird.y < 100 || this.bird.y > this.height - 100) ? -5 : 0;

        // 组合所有奖励
        reward = positionReward + velocityPenalty + pipeReward + edgePenalty;

        return [this.getState(), reward, false];
    }

    getState() {
        const closestPipe = this.pipes.find(pipe => pipe.x + pipe.width >= this.bird.x) || {
            x: this.width,
            y: this.height / 2,
            gap: 160
        };

        return [
            this.bird.y / this.height,
            this.bird.velocity / 10,
            (closestPipe.x - this.bird.x) / this.width,
            (closestPipe.y - this.bird.y) / this.height
        ];
    }

    checkCollision() {
        // Check if bird hits the ground or ceiling
        if (this.bird.y < 0 || this.bird.y > this.height) {
            return true;
        }

        // Check collision with pipes
        for (let pipe of this.pipes) {
            if (this.bird.x + this.bird.radius > pipe.x &&
                this.bird.x - this.bird.radius < pipe.x + pipe.width) {
                if (this.bird.y - this.bird.radius < pipe.y ||
                    this.bird.y + this.bird.radius > pipe.y + pipe.gap) {
                    return true;
                }
            }
        }
        return false;
    }

    spawnPipe() {
        const gap = 250;
        const minY = gap;
        const maxY = this.height - gap;
        const pipeY = Math.random() * (maxY - minY) + minY;
        
        // 在屏幕右侧生成新管道
        this.pipes.push({
            x: this.width,
            y: pipeY,
            width: 80,
            gap: gap,
            passed: false
        });
        
        // console.log('Spawned new pipe:', this.pipes[this.pipes.length - 1]); // 调试信息
    }

    render() {
        // 清空背景
        background(135, 206, 235);
        
        // 绘制管道
        fill(0, 255, 0);
        for (let pipe of this.pipes) {
            // 上管道
            rect(pipe.x, 0, pipe.width, pipe.y);
            // 下管道
            rect(pipe.x, pipe.y + pipe.gap, pipe.width, this.height - (pipe.y + pipe.gap));
        }
        
        // 绘制小鸟
        fill(255, 255, 0);
        circle(this.bird.x, this.bird.y, this.bird.radius * 2);
        
        // 绘制分数和其他信息
        fill(0);
        textSize(32);
        text(`Score: ${this.score}`, 10, 40);
        
        // 调试信息
        textSize(16);
        text(`Pipes: ${this.pipes.length}`, 10, 80);
        text(`Frame: ${this.frameCount}`, 10, 100);
    }
} 